package com.ti.demo.controller;

import com.ti.demo.tool.DateTimeTools;
import io.modelcontextprotocol.server.McpAsyncServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @description: com.ti.demo.controller
 * @Author Triagen
 * @Date 2025/10/16 19:35
 */
@RestController
public class McpServerController {

    @Autowired
    private McpAsyncServer mcpAsyncServer;

    private AtomicBoolean toolAdded = new AtomicBoolean(false);


    @GetMapping("/addTool")
    public Mono<String> addTool() {
        return Mono.fromCallable(() -> {
                    if (toolAdded.compareAndSet(false, true)) {
                        List<McpServerFeatures.AsyncToolSpecification> newTools = McpToolUtils.toAsyncToolSpecifications(ToolCallbacks.from(new DateTimeTools()));
                        for (McpServerFeatures.AsyncToolSpecification newTool : newTools) {
                            mcpAsyncServer.addTool(newTool).subscribe();
                        }
                    }
                    return "success";
                }
        );
    }
}
